import c4d
import os

from c4d import gui, plugins, bitmaps
from c4d.modules import mograph as mo

PLUGIN_ID = 1037246

NONE_ID = -1
EXECUTE_ID =1000

def SingleKeyframeCover(doc, obj, matrix):

    TRACK_POS_X = "Position . X"
    TRACK_POS_Y = "Position . Y"
    TRACK_POS_Z = "Position . Z"
    TRACK_SCALE_X = "Scale . X"
    TRACK_SCALE_Y = "Scale . Y"
    TRACK_SCALE_Z = "Scale . Z"
    TRACK_ROTATION_H = "Rotation . H"
    TRACK_ROTATION_B = "Rotation . B"
    TRACK_ROTATION_P = "Rotation . P"
    
    TRACK_KEYWORD_ARR = [TRACK_POS_X, TRACK_POS_Y, TRACK_POS_Z, TRACK_SCALE_X, TRACK_SCALE_Y, TRACK_SCALE_Z, TRACK_ROTATION_H, TRACK_ROTATION_B, TRACK_ROTATION_P]

    objTracks = obj.GetCTracks()
    if len(objTracks) == 0 : return
    
    for track in objTracks:
    
        if not track.GetName() in TRACK_KEYWORD_ARR:continue;
    
        curve = track.GetCurve()
        keyDict = curve.FindKey(doc.GetTime())
        key = None
        if keyDict != None:
            key = keyDict["key"]

        pos = matrix.off
        scale = c4d.Vector(matrix.v1.GetLength(), matrix.v2.GetLength(), matrix.v3.GetLength())
        rot = c4d.utils.MatrixToHPB(matrix)
        
        if key == None:
            key = curve.AddKey(doc.GetTime())["key"]    #Current time
    
        if track.GetName() == TRACK_POS_X:
            key.SetValue(curve, pos.x)
            
        if track.GetName() == TRACK_POS_Y:
            key.SetValue(curve, pos.y)
            
        if track.GetName() == TRACK_POS_Z:
            key.SetValue(curve, pos.z)
            
        if track.GetName() == TRACK_SCALE_X:
            key.SetValue(curve, scale.x)
            
        if track.GetName() == TRACK_SCALE_Y:
            key.SetValue(curve, scale.y)
            
        if track.GetName() == TRACK_SCALE_Z:
            key.SetValue(curve, scale.z)
            
        if track.GetName() == TRACK_ROTATION_H:
            key.SetValue(curve, rot.x)
            
        if track.GetName() == TRACK_ROTATION_P:
            key.SetValue(curve, rot.y)
            
        if track.GetName() == TRACK_ROTATION_B:
            key.SetValue(curve, rot.z)

def CreateKeyframeTrack(obj):
    trPosX = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_POSITION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_X, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trPosX) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trPosX))
    
    trPosY = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_POSITION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Y, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trPosY) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trPosY))
    
    trPosZ = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_POSITION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Z, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trPosZ) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trPosZ))
    
    trScaleX = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_SCALE, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_X, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trScaleX) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trScaleX))
    
    trScaleY = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_SCALE, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Y, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trScaleY) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trScaleY))
    
    trScaleZ = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_SCALE, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Z, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trScaleZ) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trScaleZ))
    
    trRotX = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_ROTATION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_X, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trRotX) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trRotX))
    
    trRotY = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_ROTATION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Y, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trRotY) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trRotY))
    
    trRotZ = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_ROTATION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Z, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trRotZ) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trRotZ))

def BakeAllKeyframe(doc, obj, backReferenceObj):
    CreateKeyframeTrack(obj)
    SingleKeyframeCover(doc,obj,backReferenceObj)

def Current_state_to_object(op, container=c4d.BaseContainer()):
    
    if not isinstance(op, c4d.BaseObject):
        raise TypeError('expected c4d.BaseObject, got %s.' %
                        op.__class__.__name__)

    doc = op.GetDocument()

    if not doc and (op.GetNext() or op.GetPred() or op.GetUp()):
        flags = c4d.COPYFLAGS_NO_ANIMATION | c4d.COPYFLAGS_NO_BITS | c4d.COPYFLAGS_NO_MATERIALPREVIEW
        op = op.GetClone(flags)

    if not doc:
        csto_doc = c4d.documents.BaseDocument()
        csto_doc.InsertObject(op)
    else:
        csto_doc = doc

    result = c4d.utils.SendModelingCommand(
        c4d.MCOMMAND_CURRENTSTATETOOBJECT, [op], c4d.MODELINGCOMMANDMODE_ALL,
        container, csto_doc)[0]

    if not doc:
        op.Remove()

    return result

def ExecuteMoGraphCollapse(doc, obj):
    
    collapseRoot = None
    moGraphObject = None
    moData = None
    currentCollapseObject = None
    
    moData = mo.GeGetMoData(obj)
    moGraphObject = obj
    
    collapseObject = Current_state_to_object(obj)
    collapseRoot = collapseObject
    currentCollapseObject = collapseObject.GetChildren()
    doc.InsertObject(collapseObject)

    fps = doc.GetFps()
    minTime = doc.GetMinTime()
    frameCount = minTime.GetFrame(fps)
    
    while True:

        if c4d.BaseTime(frameCount, fps) > doc.GetMaxTime():
            break

        doc.SetTime(c4d.BaseTime(frameCount, fps))
        
        executeResult = doc.ExecutePasses(None, True, True, True, c4d.BUILDFLAGS_0)
        c4d.EventAdd()

        if not executeResult:
            print "error exit"
            break

        i = 0
        moDataMatrixs = moData.GetArray(c4d.MODATA_MATRIX)
        BakeAllKeyframe(doc, collapseRoot, moGraphObject.GetMg())
        for item in currentCollapseObject:
            BakeAllKeyframe(doc, currentCollapseObject[i], moDataMatrixs[i])
            i = i + 1

        frameCount = frameCount + 1

        #------------StatusSetBar
        currentTime = c4d.BaseTime(frameCount, fps).Get() - minTime.Get()
        maxTime = doc.GetMaxTime().Get() - minTime.Get()
        
        if currentTime <= 0: currentTime = 1
        if maxTime <= 0: maxTime = 1
        
        c4d.StatusSetBar((currentTime / maxTime) * 100)
        #------------------------
    
class MyDialog(gui.GeDialog):

    def CreateLayout(self):
        self.SetTitle(title="MoGraph Collapse")
        
        self.AddButton(id=EXECUTE_ID, flags=c4d.BFH_LEFT, initw=700, name="Execute")
        
        return True

    def Command(self, id, msg):
        
        doc = c4d.documents.GetActiveDocument()
        
        if id == EXECUTE_ID:
            activeObjects = doc.GetActiveObjects(c4d.GETACTIVEOBJECTFLAGS_0)
            
            if len(activeObjects) < 1:

                gui.MessageDialog("Please Selection MoGraph Object!")

                return True
            
            i = 0
            for item in activeObjects:
            
                if mo.GeGetMoData(item) == None:
                    continue
            
                c4d.StatusSetText("Task: {0}".format(i))
                ExecuteMoGraphCollapse(doc, item)
                
                i = i + 1
                
        c4d.StatusClear()
        c4d.EventAdd()
        
        return True

class Entry(plugins.CommandData):

    dialog = None

    def Execute(self, doc):
        if self.dialog is None:
           self.dialog = MyDialog()

        return self.dialog.Open(dlgtype=c4d.DLG_TYPE_ASYNC, pluginid = PLUGIN_ID)

    def RestoreLayout(self, sec_ref):
        if self.dialog is None:
           self.dialog = MyDialog()

        return self.dialog.Restore(pluginid=PLUGIN_ID, secret = sec_ref)

if __name__ == "__main__":
    bmp = bitmaps.BaseBitmap()
    dir, file = os.path.split(__file__)
    fn = os.path.join(dir, "res", "img.png")
    bmp.InitWith(fn)
    plugins.RegisterCommandPlugin(id = PLUGIN_ID, str = "MoGraphCollapser",
                                info = 0, icon = bmp, 
                                help = "MoGraph Collapse",
                                dat = Entry())